import os
from PIL import Image
import logging
from typing import List
from datetime import datetime
from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                           QLabel, QFileDialog, QListWidget, QListWidgetItem,
                           QScrollArea, QSplitter, QApplication, QComboBox,
                           QSpinBox, QCheckBox, QGroupBox, QFormLayout,
                           QToolButton, QSizePolicy)
from PyQt6.QtCore import Qt, QSize, QPropertyAnimation, QEasingCurve, QThread, pyqtSignal
from PyQt6.QtGui import QPixmap, QImage, QIcon
from rembg import remove, new_session
import numpy as np
from config.rembg_config import RembgConfig
from common.widgets import CollapsibleBox
import concurrent.futures
from queue import Queue
import threading

class BackgroundRemover(QWidget):
    def __init__(self):
        super().__init__()
        # 设置模型下载地址为国内镜像
        os.environ['REMBG_MODEL_URL'] = "https://mirror.ghproxy.com/https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
        # 设置模型缓存目录
        os.environ['U2NET_HOME'] = os.path.join(os.path.expanduser("~"), ".u2net")
        # 创建模型缓存目
        os.makedirs(os.environ['U2NET_HOME'], exist_ok=True)
        
        self.images_path: List[str] = []
        # 设置默认输出目录
        self.output_dir = os.path.join(os.path.expanduser("~"), "Pictures", "BackgroundRemover")
        
        # 检查并创建输出目录
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir, exist_ok=True)
            logging.info(f"创建输出目录: {self.output_dir}")
        
        self.config = RembgConfig()  # 创建默认配置
        self.init_ui()
        
    def init_ui(self):
        # 创建主滚动区域
        main_scroll = QScrollArea()
        main_scroll.setWidgetResizable(True)
        main_scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
        
        # 创建主容器
        main_widget = QWidget()
        main_widget.setStyleSheet("""
            QWidget {
                background-color: #f8f9fa;
            }
            QPushButton {
                background-color: #ffffff;
                border: 1px solid #dee2e6;
                padding: 5px 15px;
                border-radius: 4px;
            }
            QPushButton:hover {
                background-color: #e9ecef;
            }
            QGroupBox {
                background-color: white;
                border: 1px solid #dee2e6;
                border-radius: 4px;
                margin-top: 12px;
            }
            QGroupBox::title {
                background-color: transparent;
            }
        """)
        
        layout = QVBoxLayout(main_widget)
        layout.setContentsMargins(20, 20, 20, 20)
        layout.setSpacing(20)
        
        # 创建分割器
        splitter = QSplitter(Qt.Orientation.Horizontal)
        
        # 左侧面板：文件列表和参数设置
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_layout.setContentsMargins(0, 0, 0, 0)
        left_layout.setSpacing(15)
        
        # 文件列表部分
        file_list_box = QGroupBox("图片列表")
        file_list_box.setStyleSheet("""
            QGroupBox {
                padding-top: 25px;
                margin-top: 15px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                padding: 5px 8px;
                font-weight: bold;
                font-size: 13px;
            }
        """)
        file_list_container = QWidget()
        file_list_layout = QVBoxLayout(file_list_container)
        file_list_layout.setSpacing(10)
        
        self.file_list = QListWidget()
        self.file_list.setIconSize(QSize(100, 100))
        self.file_list.itemSelectionChanged.connect(self.on_selection_changed)
        self.file_list.setMinimumHeight(250)
        file_list_layout.addWidget(self.file_list)
        
        # 操作按钮
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(10)
        select_btn = QPushButton("添加图片")
        select_btn.setMinimumHeight(28)
        select_btn.clicked.connect(self.select_images)
        clear_btn = QPushButton("清除全部")
        clear_btn.setMinimumHeight(28)
        clear_btn.clicked.connect(self.clear_selection)
        delete_btn = QPushButton("删除所选")
        delete_btn.setMinimumHeight(28)
        delete_btn.clicked.connect(self.delete_selected)
        btn_layout.addWidget(select_btn)
        btn_layout.addWidget(delete_btn)
        btn_layout.addWidget(clear_btn)
        file_list_layout.addLayout(btn_layout)
        
        file_list_box.setLayout(file_list_layout)
        left_layout.addWidget(file_list_box)
        
        # 参数设置部分
        settings_box = QGroupBox("抠图参数设置")
        settings_box.setStyleSheet("""
            QGroupBox {
                padding-top: 25px;
                margin-top: 15px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                padding: 5px 8px;
                font-weight: bold;
                font-size: 13px;
            }
        """)
        settings_container = QWidget()
        settings_layout = QFormLayout(settings_container)
        settings_layout.setSpacing(12)
        settings_layout.setContentsMargins(15, 10, 15, 15)
        
        # 调整所有下拉框和输入框的高度
        widget_style = """
            QComboBox, QSpinBox {
                min-height: 24px;
                padding: 2px 5px;
            }
            QCheckBox {
                min-height: 20px;
            }
        """
        settings_container.setStyleSheet(widget_style)
        
        # 预设选择
        self.preset_combo = QComboBox()
        self.preset_combo.addItems(RembgConfig.get_presets().keys())
        self.preset_combo.currentTextChanged.connect(self.load_preset)
        settings_layout.addRow("预设方案:", self.preset_combo)
        
        # 模型选择
        self.model_combo = QComboBox()
        for model, desc in RembgConfig.get_model_descriptions().items():
            self.model_combo.addItem(f"{model} - {desc}", model)
        self.model_combo.currentIndexChanged.connect(self.update_config)
        settings_layout.addRow("模型选择:", self.model_combo)
        
        # Alpha matting 设置
        self.alpha_matting_check = QCheckBox("启用")
        self.alpha_matting_check.setChecked(self.config.alpha_matting)
        self.alpha_matting_check.stateChanged.connect(self.update_config)
        settings_layout.addRow("Alpha matting:", self.alpha_matting_check)
        
        # 前景阈值
        self.fg_threshold = QSpinBox()
        self.fg_threshold.setRange(200, 255)
        self.fg_threshold.setValue(self.config.alpha_matting_foreground_threshold)
        self.fg_threshold.valueChanged.connect(self.update_config)
        settings_layout.addRow("前景阈值 (200-255):", self.fg_threshold)
        
        # 背景阈值
        self.bg_threshold = QSpinBox()
        self.bg_threshold.setRange(0, 50)
        self.bg_threshold.setValue(self.config.alpha_matting_background_threshold)
        self.bg_threshold.valueChanged.connect(self.update_config)
        settings_layout.addRow("背景阈值 (0-50):", self.bg_threshold)
        
        # 腐蚀大小
        self.erode_size = QSpinBox()
        self.erode_size.setRange(0, 30)
        self.erode_size.setValue(self.config.alpha_matting_erode_size)
        self.erode_size.valueChanged.connect(self.update_config)
        settings_layout.addRow("腐蚀大小 (0-30):", self.erode_size)
        
        # 后处理
        self.post_process_check = QCheckBox("启用")
        self.post_process_check.setChecked(self.config.post_process_mask)
        self.post_process_check.stateChanged.connect(self.update_config)
        settings_layout.addRow("后处理优化:", self.post_process_check)
        
        # 添加处理按钮到参数设置区域底部
        process_btn = QPushButton("开始处理")
        process_btn.setStyleSheet("""
            QPushButton {
                background-color: #007bff;
                color: white;
                font-weight: bold;
                padding: 8px 20px;
            }
            QPushButton:hover {
                background-color: #0056b3;
            }
        """)
        process_btn.clicked.connect(self.process_images)
        settings_layout.addRow("", process_btn)
        
        settings_box.setLayout(settings_layout)
        left_layout.addWidget(settings_box)
        
        # 输出目录选择
        output_box = QGroupBox("输出设置")
        output_box.setStyleSheet("""
            QGroupBox {
                padding-top: 25px;
                margin-top: 15px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                padding: 5px 8px;
                font-weight: bold;
                font-size: 13px;
            }
        """)
        output_layout = QHBoxLayout()
        output_layout.setContentsMargins(15, 10, 15, 15)
        output_layout.setSpacing(10)
        self.output_label = QLabel(f"输出目录：{self.output_dir}")
        self.output_label.setWordWrap(True)
        output_btn = QPushButton("选择输出目录")
        output_btn.setMinimumHeight(32)
        output_btn.clicked.connect(self.select_output_dir)
        output_layout.addWidget(output_btn)
        output_layout.addWidget(self.output_label)
        output_layout.addStretch()
        output_box.setLayout(output_layout)
        left_layout.addWidget(output_box)
        
        # 右侧面板：预览区域
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        right_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.setSpacing(15)
        
        # 预览区域
        preview_group = QGroupBox("图片预览")
        preview_group.setStyleSheet("""
            QGroupBox {
                padding-top: 25px;
                margin-top: 15px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                padding: 5px 8px;
                font-weight: bold;
                font-size: 13px;
            }
        """)
        preview_layout = QVBoxLayout()
        preview_layout.setContentsMargins(15, 10, 15, 15)
        preview_layout.setSpacing(15)
        
        # 原图预览
        original_label = QLabel("原图:")
        preview_layout.addWidget(original_label)
        
        self.original_preview = QLabel()
        self.original_preview.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.original_preview.setFixedSize(400, 250)
        self.original_preview.setStyleSheet("border: 1px solid #dee2e6;")
        preview_layout.addWidget(self.original_preview)
        
        # 抠图预览
        preview_header = QHBoxLayout()
        result_label = QLabel("抠图预览:")
        preview_btn = QPushButton("预览效果")
        preview_btn.clicked.connect(self.preview_selected)
        preview_header.addWidget(result_label)
        preview_header.addWidget(preview_btn)
        preview_header.addStretch()
        preview_layout.addLayout(preview_header)
        
        self.result_preview = QLabel()
        self.result_preview.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.result_preview.setFixedSize(400, 250)
        self.result_preview.setStyleSheet("border: 1px solid #dee2e6;")
        preview_layout.addWidget(self.result_preview)
        
        preview_group.setLayout(preview_layout)
        right_layout.addWidget(preview_group)
        right_layout.addStretch()
        
        # 添加面板到分割器
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([300, 700])
        
        # 添加分割器到主布局
        layout.addWidget(splitter)
        
        # 进度显示框
        progress_frame = QWidget()
        progress_layout = QVBoxLayout(progress_frame)
        progress_layout.setContentsMargins(0, 0, 0, 0)
        
        self.progress_text = QLabel()
        self.progress_text.setWordWrap(True)
        self.progress_text.setStyleSheet("""
            QLabel {
                background-color: white;
                border: 1px solid #dee2e6;
                border-radius: 4px;
                padding: 15px;
                font-size: 12px;
                line-height: 1.5;
            }
        """)
        self.progress_text.setMinimumHeight(120)
        progress_layout.addWidget(self.progress_text)
        
        layout.addWidget(progress_frame)
        
        # 设置主滚动区域的widget
        main_scroll.setWidget(main_widget)
        
        # 设置主布局
        main_layout = QVBoxLayout(self)
        main_layout.addWidget(main_scroll)
    
    def select_images(self):
        """选择图片文件"""
        files, _ = QFileDialog.getOpenFileNames(
            self,
            "选择图片",
            "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp);;所有文件 (*.*)"
        )
        if files:
            self.images_path.extend(files)
            self.update_file_list()
    
    def update_file_list(self):
        """更新文件列表显示"""
        self.file_list.clear()
        for path in self.images_path:
            item = QListWidgetItem()
            pixmap = QPixmap(path)
            pixmap = pixmap.scaled(100, 100, Qt.AspectRatioMode.KeepAspectRatio)
            icon = QIcon(pixmap)
            item.setIcon(icon)
            item.setText(os.path.basename(path))
            item.setData(Qt.ItemDataRole.UserRole, path)
            self.file_list.addItem(item)
    
    def process_image(self, img, preview=False):
        """使用rembg处理图片"""
        try:
            # 使用当前配置创建会话
            session = new_session(
                model_name=self.config.model_name,
                providers=self.config.providers,
                parameters={
                    "alpha_matting": self.config.alpha_matting,
                    "alpha_matting_foreground_threshold": self.config.alpha_matting_foreground_threshold,
                    "alpha_matting_background_threshold": self.config.alpha_matting_background_threshold,
                    "alpha_matting_erode_size": self.config.alpha_matting_erode_size,
                    "post_process_mask": self.config.post_process_mask,
                }
            )
            
            # 确保图片是RGB模式
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            # 使用rembg移除背景
            output = remove(
                img,
                session=session,
                post_process_mask=self.config.post_process_mask,
            )
            
            if preview:
                # 转换为QPixmap用于预览
                output_array = np.array(output)
                height, width, channel = output_array.shape
                bytes_per_line = channel * width
                image = QImage(output_array.data, width, height, bytes_per_line, QImage.Format.Format_RGBA8888)
                pixmap = QPixmap.fromImage(image)
                scaled_pixmap = pixmap.scaled(400, 300, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation)
                self.result_preview.setPixmap(scaled_pixmap)
                QApplication.processEvents()
            
            return output
            
        except Exception as e:
            logging.error(f"处理图片失败: {str(e)}")
            raise e

    def on_selection_changed(self):
        """处理选择变化事件"""
        items = self.file_list.selectedItems()
        if items:
            path = items[0].data(Qt.ItemDataRole.UserRole)
            # 只显示原图预览
            pixmap = QPixmap(path)
            scaled_pixmap = pixmap.scaled(
                400, 300,
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.original_preview.setPixmap(scaled_pixmap)
            self.result_preview.clear()  # 清除之前的抠图预览
        else:
            self.original_preview.clear()
            self.result_preview.clear()
    
    def delete_selected(self):
        """删除选中的图片"""
        items = self.file_list.selectedItems()
        for item in items:
            path = item.data(Qt.ItemDataRole.UserRole)
            self.images_path.remove(path)
            self.file_list.takeItem(self.file_list.row(item))
        self.original_preview.clear()
    
    def clear_selection(self):
        """清除所有选择"""
        self.images_path.clear()
        self.file_list.clear()
        self.original_preview.clear()
        self.result_preview.clear()  # 清除预览
        self.progress_text.clear()   # 清除进度信息
    
    def select_output_dir(self):
        """选择输出目录"""
        dir_path = QFileDialog.getExistingDirectory(
            self,
            "选择输出目录",
            self.output_dir
        )
        if dir_path:
            self.output_dir = dir_path
            self.output_label.setText(f"输出目录：{dir_path}")
    
    def process_images(self):
        """处理所有图片"""
        if not self.images_path:
            msg = "请先选择图片"
            print(msg)
            self.progress_text.setText(msg)
            logging.info("未选择图片，操作终止")
            return
        
        try:
            total = len(self.images_path)
            start_msg = f"开始处理图片，共 {total} 张"
            print("\n" + start_msg)
            self.progress_text.setText(start_msg + "\n")
            logging.info(start_msg)
            
            # 创建输出目录
            output_dir = os.path.join(
                self.output_dir,
                f"removed_bg_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            )
            os.makedirs(output_dir, exist_ok=True)
            dir_msg = f"创建输出目录: {output_dir}"
            print(dir_msg)
            logging.info(dir_msg)
            
            # 创建��启动工作线程
            self.progress_text.clear()
            self.worker = ImageProcessWorker(self.images_path, output_dir, self.config)
            self.worker.progress.connect(self.update_progress)
            self.worker.finished.connect(self.process_complete)
            self.worker.start()
            
        except Exception as e:
            error_msg = f"处理失败: {str(e)}"
            print("\n" + error_msg)
            logging.error(error_msg, exc_info=True)
            self.progress_text.setText(error_msg)
    
    def open_output_dir(self, path: str):
        """打开输出目录"""
        try:
            os.startfile(path)
        except Exception as e:
            logging.error(f"打开输出目录失败: {str(e)}")
            self.info_label.setText(f"打开输出目录失败: {str(e)}") 
    
    def load_preset(self, preset_name: str):
        """加载预设参数"""
        if preset_name in RembgConfig.get_presets():
            preset = RembgConfig.get_presets()[preset_name]
            self.config = RembgConfig(**preset)
            self.update_ui_from_config()
    
    def update_ui_from_config(self):
        """根据配置更新UI控件"""
        # 更新UI控件的值，但不触发update_config
        self.model_combo.setCurrentText(self.config.model_name)
        self.alpha_matting_check.setChecked(self.config.alpha_matting)
        self.fg_threshold.setValue(self.config.alpha_matting_foreground_threshold)
        self.bg_threshold.setValue(self.config.alpha_matting_background_threshold)
        self.erode_size.setValue(self.config.alpha_matting_erode_size)
        self.post_process_check.setChecked(self.config.post_process_mask)
    
    def update_config(self):
        """从UI更新配置"""
        self.config.model_name = self.model_combo.currentData()
        self.config.alpha_matting = self.alpha_matting_check.isChecked()
        self.config.alpha_matting_foreground_threshold = self.fg_threshold.value()
        self.config.alpha_matting_background_threshold = self.bg_threshold.value()
        self.config.alpha_matting_erode_size = self.erode_size.value()
        self.config.post_process_mask = self.post_process_check.isChecked() 
    
    def preview_selected(self):
        """预览选中图片的抠图效果"""
        items = self.file_list.selectedItems()
        if not items:
            return
        
        try:
            path = items[0].data(Qt.ItemDataRole.UserRole)
            with Image.open(path) as img:
                self.process_image(img, preview=True)
        except Exception as e:
            logging.error(f"预览失败: {str(e)}")
            self.result_preview.clear() 

    def update_progress(self, message, is_progress):
        """更新进度信息"""
        current_text = self.progress_text.text()
        if current_text:
            current_text += "\n"
        self.progress_text.setText(current_text + message)
        print(message)
        if is_progress:
            logging.info(message)
        else:
            logging.error(message)

    def process_complete(self, stats):
        """处理完成的回调"""
        complete_msg = (
            f"\n处理完成！\n"
            f"总数: {stats['total']} 张\n"
            f"成功: {stats['success']} 张\n"
            f"失败: {stats['failed']} 张\n"
            f"成功率: {(stats['success']/stats['total'])*100:.1f}%\n"
            f"输出目录: {stats['output_dir']}"
        )
        
        # 更新界面显示
        current_text = self.progress_text.text()
        if current_text:
            current_text += "\n"
        self.progress_text.setText(current_text + complete_msg)
        print("\n" + complete_msg)
        logging.info(complete_msg)
        
        # 添加打开输出目录按钮
        open_dir_btn = QPushButton("打开输出目录")
        open_dir_btn.clicked.connect(lambda: self.open_output_dir(stats['output_dir']))
        open_dir_btn.setStyleSheet("""
            QPushButton {
                background-color: #28a745;
                color: white;
                font-weight: bold;
                padding: 8px 20px;
            }
            QPushButton:hover {
                background-color: #218838;
            }
        """)
        
        # 更新按钮
        for i in reversed(range(self.layout().count())):
            widget = self.layout().itemAt(i).widget()
            if isinstance(widget, QPushButton) and widget.text() == "打开输出目录":
                widget.deleteLater()
        
        self.layout().addWidget(open_dir_btn)

class ImageProcessWorker(QThread):
    """图片处理工作线程"""
    progress = pyqtSignal(str, bool)  # 发送进度信息和是否成功
    finished = pyqtSignal(dict)  # 发送处理完成的统计信息
    
    def __init__(self, images, output_dir, config):
        super().__init__()
        self.images = images
        self.output_dir = output_dir
        self.config = config
        self.total = len(images)
        self.processed = 0
        self.success = 0
        self.failed = 0
    
    def run(self):
        # 创建线程池
        with concurrent.futures.ThreadPoolExecutor(max_workers=min(4, len(self.images))) as executor:
            # 提交所有任务
            futures = {executor.submit(self.process_single_image, img_path): img_path 
                      for img_path in self.images}
            
            # 处理完成的任务
            for future in concurrent.futures.as_completed(futures):
                img_path = futures[future]
                try:
                    success = future.result()
                    self.processed += 1
                    if success:
                        self.success += 1
                    else:
                        self.failed += 1
                except Exception as e:
                    self.processed += 1
                    self.failed += 1
                    self.progress.emit(f"✗ {os.path.basename(img_path)} 处理失败: {str(e)}", False)
        
        # 发送完成信息
        self.finished.emit({
            'total': self.total,
            'success': self.success,
            'failed': self.failed,
            'output_dir': self.output_dir
        })
    
    def process_single_image(self, img_path):
        try:
            current_file = os.path.basename(img_path)
            self.progress.emit(f"正在处理 ({self.processed + 1}/{self.total}) "
                             f"[{((self.processed + 1) / self.total) * 100:.1f}%]: {current_file}", True)
            
            with Image.open(img_path) as img:
                # 创建会话
                session = new_session(
                    model_name=self.config.model_name,
                    providers=self.config.providers,
                    parameters={
                        "alpha_matting": self.config.alpha_matting,
                        "alpha_matting_foreground_threshold": self.config.alpha_matting_foreground_threshold,
                        "alpha_matting_background_threshold": self.config.alpha_matting_background_threshold,
                        "alpha_matting_erode_size": self.config.alpha_matting_erode_size,
                        "post_process_mask": self.config.post_process_mask,
                    }
                )
                
                # 确保图片是RGB模式
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # 处理图片
                result = remove(
                    img,
                    session=session,
                    post_process_mask=self.config.post_process_mask,
                )
                
                # 保存结果
                output_filename = os.path.splitext(current_file)[0] + '.png'
                save_path = os.path.join(self.output_dir, f"nobg_{output_filename}")
                result.save(save_path, "PNG")
                
                self.progress.emit(f"✓ {current_file} 处理成功", True)
                return True
                
        except Exception as e:
            self.progress.emit(f"✗ {current_file} 处理失败: {str(e)}", False)
            return False